{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import sys\n",
    "from simulator import (\n",
    "    Simulator, string_to_micro, micro_to_time,\n",
    "    BUY, SELL, SHORT, EXCH_INET,\n",
    "    BOOK_DEPTH1_PRICE, ORDER_EVENTS,\n",
    "    )\n",
    "import numpy as np\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "from sklearn.linear_model import LinearRegression\n",
    "from statsmodels.tsa.ar_model import AR\n",
    "from scipy import integrate \n",
    "from scipy import optimize\n",
    "\n",
    "\n",
    "class Ave_Lee(object):\n",
    "    def __init__(self, session, date, tickers, start_time, end_time, pnl, buy_shares1, buy_dollars, \n",
    "                 sell_shares1, sell_dollars, buy_shares2, sell_shares2, threshold, interval, training_size, trades):\n",
    "        self.session = session\n",
    "        self.date = date\n",
    "        self.tickers = tickers\n",
    "        self.ticker1 = self.tickers[0]\n",
    "        self.ticker2 = self.tickers[1]\n",
    "        self.start_time = start_time\n",
    "        self.end_time = end_time\n",
    "        self.halt_trading = string_to_micro('5m') # close position 5 minutes before trading ends\n",
    "        self.interval = string_to_micro(interval)\n",
    "        self.interval_pnl = string_to_micro('1s')\n",
    "        self.state = 'NULL'\n",
    "        \n",
    "        # variables for BUY or SELL\n",
    "        self.side1 = 0\n",
    "        self.side2 = 0\n",
    "        \n",
    "        # variables for order size\n",
    "        self.order_size1 = 100\n",
    "        self.order_size2 = 1\n",
    "        self.buy_size = 10000 * 10**6 # how much we can buy for in each trade\n",
    "        \n",
    "        # variables to keep track of total shares bought/sold and the corresponding amount of money\n",
    "        self.buy_shares1 = buy_shares1\n",
    "        self.buy_dollars = buy_dollars\n",
    "        self.sell_shares1 = sell_shares1\n",
    "        self.sell_dollars = sell_dollars\n",
    "        self.buy_shares2 = buy_shares2\n",
    "        self.sell_shares2 = sell_shares2\n",
    "        \n",
    "        # variables to keep track of how many positions we have opened and closed respectively\n",
    "        self.trades = trades\n",
    "        \n",
    "        # variables used for the fitOU, when to open/close a position and how far we look back\n",
    "        self.dt = 1\n",
    "        self.long_open = threshold[0]\n",
    "        self.long_close = threshold[1]\n",
    "        self.short_open = threshold[0]\n",
    "        self.short_close = threshold[2]\n",
    "        self.training_size = training_size\n",
    "        \n",
    "        # Start timer/call the start_callback function\n",
    "        self.session.add_timer(self.start_time, self.start_callback)\n",
    "        \n",
    "        # List to store pnl every time we update it\n",
    "        self.pnl = pnl\n",
    "        \n",
    "        # Dictionary to store time, midprices and the returns each timestep\n",
    "        self.results = {'time': []}\n",
    "        for ticker in self.tickers:\n",
    "            self.results[ticker] = []\n",
    "            self.results['return {}'.format(ticker)] = []\n",
    "        \n",
    "        # subscribe to the tickers of interest, and set the first timer\n",
    "        for ticker in self.tickers:\n",
    "            self.session.subscribe_ticker_all_feeds(ticker)\n",
    "    \n",
    "    \n",
    "    def start_callback(self, time):\n",
    "        # Start subscribing to executed orders which then calls event_callback\n",
    "        for ticker in self.tickers:\n",
    "            self.session.subscribe_event(ticker, ORDER_EVENTS, self.event_callback)\n",
    "        # Add timers for updating pnl and running the algo\n",
    "        self.session.add_timer(time, self.timer_callback_pnl)\n",
    "        self.session.add_timer(time, self.timer_callback_algo)\n",
    "        \n",
    "    \n",
    "    def event_callback(self, ticker, event_params):\n",
    "        # Call the execution manager whenever we have an execution\n",
    "        self.process_executions(event_params)\n",
    "        \n",
    "        \n",
    "    def timer_callback_pnl(self, time):\n",
    "        # This timer callback updates pnl every second to see how it evolves over the day            \n",
    "        pnl = self.get_pnl()\n",
    "        self.pnl.append(pnl / 1000000.0)\n",
    "        \n",
    "        if time < self.end_time:\n",
    "            self.session.add_timer(time + self.interval_pnl, self.timer_callback_pnl)\n",
    "            \n",
    "    \n",
    "    def timer_callback_algo(self, time):\n",
    "        # This timer callback performs the algorithm\n",
    "        \n",
    "        # Append the time\n",
    "        self.results['time'].append(micro_to_time(time))       \n",
    "        # Append the midprices\n",
    "        self.results[self.ticker1].append(self.get_midmarket(self.ticker1) / 1000000.0)\n",
    "        self.results[self.ticker2].append(self.get_midmarket(self.ticker2) / 1000000.0)\n",
    "        \n",
    "        # Start calculating returns after 1 second\n",
    "        if time > self.start_time + 10**6:\n",
    "            self.results['return {}'.format(self.ticker1)].append(np.float(returns(self.results[self.ticker1][-2:])))\n",
    "            self.results['return {}'.format(self.ticker2)].append(np.float(returns(self.results[self.ticker2][-2:])))\n",
    "        \n",
    "        # Get the best bid and offer, compute the midmarket\n",
    "        bid1, ask1 = self.session.get_inside_market(self.ticker1)\n",
    "        bid2, ask2 = self.session.get_inside_market(self.ticker2)\n",
    "        \n",
    "        # Start collecting signals after training_size * 1 second + halt trading\n",
    "        if time > self.start_time + self.halt_trading + self.training_size * 10**6:\n",
    "            # Collect the last training_size of returns\n",
    "            returns1 = self.results['return {}'.format(self.ticker1)][-self.training_size:]\n",
    "            returns2 = self.results['return {}'.format(self.ticker2)][-self.training_size:]\n",
    "            \n",
    "            # Regress the returns and fit the residuals, calculate the s-score\n",
    "            residuals, a,b = regress(returns1,returns2)\n",
    "            kappa, m, sigma, sigmaeq = fitOU(residuals, self.training_size)\n",
    "            s = sscore(m, sigmaeq)\n",
    "            \n",
    "            # Find current net position (=0: neutral, <0: we are short asset 1, >0: we are long asset 1)\n",
    "            pos = self.buy_shares1 - self.sell_shares1\n",
    "            \n",
    "            # Feature to check if we have orders at the market before we open a new position\n",
    "            orders = self.session.get_all_orders()        \n",
    "            \n",
    "            if not orders and b > 0 and kappa > 0 and sigma > 0 and time < self.end_time - self.halt_trading:\n",
    "                if pos == 0:\n",
    "                    if s < -self.long_open:\n",
    "                        self.side1 = BUY\n",
    "                        self.side2 = SELL\n",
    "                        price1 = ask1['price']\n",
    "                        price2 = bid2['price']\n",
    "                        \n",
    "                        # The regression coefficient determines the sell size for ticker 2\n",
    "                        self.order_size1 = self.buy_size // price1\n",
    "                        self.order_size2 = int(b * self.order_size1)\n",
    "                        self.session.add_order(self.ticker1, self.side1, self.order_size1, price1, exchange=EXCH_INET)\n",
    "                        self.session.add_order(self.ticker2, self.side2, self.order_size2, price2, exchange=EXCH_INET)\n",
    "                        self.trades += 1                   \n",
    "                    elif s > self.short_open:\n",
    "                        self.side1 = SELL\n",
    "                        self.side2 = BUY\n",
    "                        price1 = bid1['price']\n",
    "                        price2 = ask2['price']\n",
    "                        \n",
    "                        # The regression coefficient determines the buy size for ticker 2 \n",
    "                        self.order_size1 = self.buy_size // price1\n",
    "                        self.order_size2 = int(b * self.order_size1)\n",
    "                        \n",
    "                        self.session.add_order(self.ticker1, self.side1, self.order_size1, price1, exchange=EXCH_INET)\n",
    "                        self.session.add_order(self.ticker2, self.side2, self.order_size2, price2, exchange=EXCH_INET)\n",
    "                        self.trades += 1                    \n",
    "                elif pos < 0 and s < self.short_close:\n",
    "                    self.side1 = BUY\n",
    "                    self.side2 = SELL\n",
    "                    \n",
    "                    price1 = ask1['price']\n",
    "                    price2 = bid2['price']\n",
    "                    \n",
    "                    self.session.add_order(self.ticker1, self.side1, self.order_size1, price1, exchange=EXCH_INET)\n",
    "                    self.session.add_order(self.ticker2, self.side2, self.order_size2, price2, exchange=EXCH_INET)  \n",
    "                elif pos > 0 and s > -self.long_close:\n",
    "                    self.side1 = SELL\n",
    "                    self.side2 = BUY\n",
    "                    \n",
    "                    price1 = bid1['price']\n",
    "                    price2 = ask2['price']\n",
    "                    \n",
    "                    self.session.add_order(self.ticker1, self.side1, self.order_size1, price1, exchange=EXCH_INET)\n",
    "                    self.session.add_order(self.ticker2, self.side2, self.order_size2, price2, exchange=EXCH_INET)\n",
    "        \n",
    "        # Close any open position halt_trading time steps before trading ends\n",
    "        if time >= self.end_time - self.halt_trading and pos != 0:\n",
    "            if pos < 0:\n",
    "                self.side1 = BUY\n",
    "                self.side2 = SELL\n",
    "                \n",
    "                price1 = ask1['price']\n",
    "                price2 = bid2['price']\n",
    "                \n",
    "                self.session.add_order(self.ticker1, self.side1, self.order_size1, price1, exchange=EXCH_INET)\n",
    "                self.session.add_order(self.ticker2, self.side2, self.order_size2, price2, exchange=EXCH_INET)\n",
    "                \n",
    "            elif pos > 0:\n",
    "                self.side1 = SELL\n",
    "                self.side2 = BUY\n",
    "                \n",
    "                price1 = bid1['price']\n",
    "                price2 = ask2['price']\n",
    "                \n",
    "                self.session.add_order(self.ticker1, self.side1, self.order_size1, price1, exchange=EXCH_INET)\n",
    "                self.session.add_order(self.ticker2, self.side2, self.order_size2, price2, exchange=EXCH_INET)\n",
    "                \n",
    "        # Reset the timer unless we are done \n",
    "        if time < self.end_time:\n",
    "            self.session.add_timer(time + self.interval, self.timer_callback_algo)\n",
    "                \n",
    "            \n",
    "    def process_executions(self, evp):\n",
    "        # Make sure that we only update if we have executed any orders\n",
    "        if 'executed_orders' in evp:\n",
    "            time = self.session.current_time()\n",
    "            for ex in evp['executed_orders']:\n",
    "                order = ex['order']\n",
    "                side = order['side']\n",
    "                ticker = order['ticker']\n",
    "                \n",
    "                if ticker == self.ticker1:\n",
    "                    if side == 'B':\n",
    "                        self.buy_shares1 += ex['quantity_executed']\n",
    "                        #self.buy_dollars += ex['quantity_executed'] * ex['price_executed']\n",
    "                        \n",
    "                        # buy in midmarker to check if spread is \"eating\" profits\n",
    "                        self.buy_dollars += ex['quantity_executed'] * self.get_midmarket(ticker)\n",
    "                        \n",
    "                    else:\n",
    "                        self.sell_shares1 += ex['quantity_executed']\n",
    "                        #self.sell_dollars += ex['quantity_executed'] * ex['price_executed']\n",
    "                        \n",
    "                        # sell in midmarker to check if spread is \"eating\" profits\n",
    "                        self.sell_dollars += ex['quantity_executed'] * self.get_midmarket(ticker)\n",
    "                        \n",
    "                    pos = self.buy_shares1 - self.sell_shares1\n",
    "                elif ticker == self.ticker2:\n",
    "                    if side == 'B':\n",
    "                        self.buy_shares2 += ex['quantity_executed']\n",
    "                        #self.buy_dollars += ex['quantity_executed'] * ex['price_executed']\n",
    "                        \n",
    "                        # buy in midmarker to check if spread is \"eating\" profits\n",
    "                        self.buy_dollars += ex['quantity_executed'] * self.get_midmarket(ticker)\n",
    "                    else:\n",
    "                        self.sell_shares2 += ex['quantity_executed']\n",
    "                        #self.sell_dollars += ex['quantity_executed'] * ex['price_executed']\n",
    "                        \n",
    "                        # sell in midmarker to check if spread is \"eating\" profits\n",
    "                        self.sell_dollars += ex['quantity_executed'] * self.get_midmarket(ticker)\n",
    "                        \n",
    "                    pos = self.buy_shares2 - self.sell_shares2\n",
    "                    \n",
    "                pnl = self.get_pnl()\n",
    "                \n",
    "                \n",
    "\n",
    "    def get_midmarket(self, ticker):\n",
    "        bid, ask = self.session.get_inside_market(ticker)\n",
    "        return (bid['price'] + ask['price']) / 2.0\n",
    "    \n",
    "    \n",
    "    def get_pnl(self):\n",
    "        # Mark to the midmarket\n",
    "        mid1 = self.get_midmarket(self.ticker1)\n",
    "        mid2 = self.get_midmarket(self.ticker2)\n",
    "        pnl = self.sell_dollars - self.buy_dollars + (self.buy_shares1 - self.sell_shares1) * mid1 + (self.buy_shares2 - self.sell_shares2) * mid2\n",
    "        return pnl\n",
    "    \n",
    "    \n",
    "    def end(self):\n",
    "        return (self.pnl, self.buy_shares1, self.buy_dollars, self.sell_shares1, self.sell_dollars, self.buy_shares2, self.sell_shares2, self.trades)\n",
    "\n",
    "\n",
    "def regress(returns1,returns2):\n",
    "    x = np.asarray(returns1).reshape(-1,1)\n",
    "    y = np.asarray(returns2).reshape(-1,1)\n",
    "    model = LinearRegression()\n",
    "    model.fit(x,y)\n",
    "    a = model.intercept_[0]\n",
    "    b = model.coef_[0,0]\n",
    "    residuals = y-model.predict(x)\n",
    "    return residuals, a,b\n",
    "\n",
    "\n",
    "def returns(midprices):\n",
    "    log_return = np.diff(np.log(midprices), axis=-1)\n",
    "    return log_return\n",
    "\n",
    "\n",
    "def fitOU(residual, training_size):\n",
    "    dt = 1\n",
    "    ou = np.cumsum(residual)\n",
    "    model = AR(ou)\n",
    "    fittedmodel = model.fit(maxlag=1, disp=-1)  \n",
    "    a = fittedmodel.params[0]\n",
    "    b = fittedmodel.params[1]\n",
    "    var =  fittedmodel.sigma2\n",
    "    if b > 0.0 and b < np.exp(-2.0/training_size):\n",
    "        kappa = -np.log(b) / dt    \n",
    "        m = a / (1.0 - np.exp(-kappa * dt))\n",
    "        sigma = np.sqrt(var * 2.0 * kappa / (1.0 - np.exp(-2.0 * kappa * dt)))\n",
    "        sigmaeq = np.sqrt(var / (1.0 - np.exp(-2.0 * kappa * dt)));\n",
    "        return kappa, m, sigma, sigmaeq\n",
    "    else:\n",
    "        return -1.0,0,0,0\n",
    "\n",
    "\n",
    "def sscore(m, sigmaeq):\n",
    "    if sigmaeq != 0:\n",
    "        return -m/sigmaeq\n",
    "    elif m>0:\n",
    "        return 10000000\n",
    "    else:\n",
    "        return -10000000\n",
    "\n",
    "\n",
    "def plot_graph(pnl, ticker1, ticker2):\n",
    "    pnl_array = np.asarray(pnl)\n",
    "    plt.plot(pnl_array)\n",
    "    plt.title('PnL for {} vs {}'.format(ticker1, ticker2))\n",
    "    plt.show()\n",
    "    return None\n",
    "    \n",
    "\n",
    "# Simulate the algorithm on several pairs and for multiple days\n",
    "start_time = string_to_micro(\"9:30\")\n",
    "end_time = string_to_micro(\"16:00\")\n",
    "\n",
    "# dates to be tested on\n",
    "dates = ['20170508', '20170509', '20170510', '20170511', '20170512', '20170515', '20170516', '20170517',\n",
    "                    '20170518', '20170519']\n",
    "\n",
    "tickers1 = ['ORCL', 'AAPL', 'NVDA', 'TXN', 'AAPL']\n",
    "tickers2 = ['NVDA', 'ORCL', 'MSFT', 'AAPL', 'CSCO']\n",
    "\n",
    "# hyperparameters\n",
    "threshold = [1.4, 0.5, 0.5]\n",
    "interval = '1m' '15s'\n",
    "training_size = 250\n",
    "\n",
    "for i in range(len(tickers1)):  \n",
    "    ticker1 = tickers1[i]\n",
    "    ticker2 = tickers2[i]\n",
    "        \n",
    "    print(str(ticker1) + ' vs ' + str(ticker2))\n",
    "    \n",
    "    # Reset the pnl, and other variables when we have a new pair of tickers\n",
    "    pnl = []          \n",
    "    buy_shares1 = 0\n",
    "    buy_dollars = 0\n",
    "    sell_shares1 = 0\n",
    "    sell_dollars = 0\n",
    "    buy_shares2 = 0\n",
    "    sell_shares2 = 0       \n",
    "    trades = 0\n",
    "    \n",
    "    # Run simulator for all dates and aggregate the pnl\n",
    "    for date in dates:\n",
    "        sim = Simulator(Ave_Lee)\n",
    "        (pnl, buy_shares1, buy_dollars, sell_shares1, sell_dollars, buy_shares2, sell_shares2, trades) = sim.run(date, [ticker1, ticker2], use_om=True, start_time=start_time, end_time=end_time, \n",
    "                    pnl=pnl, buy_shares1=buy_shares1, buy_dollars=buy_dollars, sell_shares1=sell_shares1,\n",
    "                    sell_dollars=sell_dollars, buy_shares2=buy_shares2, sell_shares2=sell_shares2, threshold=threshold,\n",
    "                    interval=interval, training_size=training_size, trades=trades)\n",
    "        \n",
    "    pnl_array = np.asarray(pnl)\n",
    "         \n",
    "    print('Last PnL of {}: {}'.format(opt, pnl[-1]))\n",
    "    print('Dollar per trade for {}: {}'.format(opt, pnl[-1] / (10000 * trades)))\n",
    "        \n",
    "    plt.plot(pnl_array)\n",
    "    plt.title('PnL for {} vs {}'.format(ticker1, ticker2))\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python2",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
